#include "Equation2D.h"

using namespace std;

double f2D(double x, double y)
{
	return exp(sin(x) * sin(y)) * (2 * sin(x) * sin(y) - 
		cos(x) * cos(x) * sin(y) * sin(y) - cos(y) * cos(y) * sin(x) * sin(x));
	//return 2 * (x - 1) * (y - 1) * sin(x) * sin(y)
	//	- 2 * (y - 1) * cos(x) * sin(y) - 2 * (x - 1) * sin(x) * cos(y);
}

double u2D(double x, double y)
{
	return exp(sin(x) * sin(y));
	//return (x - 1) * sin(x) * (y - 1) * sin(y);
}

const Vector GetBoundary(int n, std::function<double(double x, double y)> _u)
{
	Vector v(4 * n);
	for (int k = 0; k <= n - 1; ++k)
		v[k] = _u(1.0 * k / n, 0.0);
	for (int k = n; k <= 2 * n - 1; ++k)
		v[k] = _u(1.0, 1.0 * (k - n) / n);
	for (int k = 2 * n; k <= 3 * n - 1; ++k)
		v[k] = _u(1.0 * (3.0 * n - k) / n, 1.0);
	for (int k = 3 * n; k <= 4 * n - 1; ++k)
		v[k] = _u(0.0, 1.0 * (4.0 * n - k) / n);
	return v;
}

SymmetricTridiagonal2D::SymmetricTridiagonal2D(int _n) : n(_n), diag(0), alpha(0) {}

SymmetricTridiagonal2D::SymmetricTridiagonal2D(int _n, double d, double a) : n(_n), diag(d), alpha(a) {}

const Vector SymmetricTridiagonal2D::operator*(Vector& x) const
{
	
	if (x.size() != (n-1) * (n-1))
		throw std::length_error("SymmetricTridiagonal: sizeof(x) != (n-1)^2");
	Vector y((n - 1) * (n - 1));
	if (n <= 2)
		y[0] = x[0] * diag;
	else {
		SymmetricTridiagonal T(n, diag, alpha);
		Vector u0 = Vector(x[slice(0, n - 1, 1)]), u1 = Vector(x[slice(n-1, n - 1, 1)]), u2;
		y[slice(0, n - 1, 1)] = T * u0 + alpha * u1;
		u0 = Vector(x[slice((n - 1) * (n - 3), n - 1, 1)]);
		u1 = Vector(x[slice((n - 1) * (n - 2), n - 1, 1)]);
		y[slice((n - 1) * (n - 2), n - 1, 1)] = alpha * u0 + T * u1;
		for (int i = 0; i < n - 3; ++i) {
			u0 = Vector(x[slice(i * (n - 1), n - 1, 1)]);
			u1 = Vector(x[slice((i + 1) * (n - 1), n - 1, 1)]);
			u2 = Vector(x[slice((i + 2) * (n - 1), n - 1, 1)]);
			y[slice((i + 1) * (n - 1), n - 1, 1)] = alpha * u0 + T * u1 + alpha * u2;
		}
	}
	return y;
}

double SymmetricTridiagonal2D::GetDiag() const
{
	return diag;
}

double SymmetricTridiagonal2D::GetAlpha() const
{
	return alpha;
}


PossionMatrix2D::PossionMatrix2D(int _n) : SymmetricTridiagonal2D(_n, 4.0 * _n * _n, -1.0 * _n * _n) {}

RHS2D::RHS2D(int _n, const Vector& boundary, std::function<double(double x, double y)> func) : n(_n)
{

	if (boundary.size() != static_cast<size_t>(4 * n))
		throw std::length_error("RHS2D(): boundary.size != 4n");
	if (n < 2)
		throw std::length_error("RHS2D(): n < 2");
	v = Vector((n - 1) * (n - 1));
	for (int i = 1; i <= n - 1; ++i)
		for (int j = 1; j <=  n - 1; ++j)
			v[(n-1)*(j-1)+(i-1)] = func(1.0 * i / n, 1.0 * j / n);
	for (int i = 1; i <= n - 1; ++i)
		v[i - 1] += boundary[i] * n * n;
	for (int j = 1; j <= n - 1; ++j)
		v[(n-1)*(j-1)+(n-2)] += boundary[n + j] * n * n;
	for (int i = n - 1; i >= 1; --i)
		v[(n-1)*(n-2)+i-1] += boundary[3 * n - i] * n * n;
	for (int j = n - 1; j >= 1; --j)
		v[(n-1)*(j-1)] += boundary[4 * n - j] * n * n;
}

const Vector& RHS2D::GetVector() const
{
	return v;
}

const int RHS2D::GetSize() const
{
	return n;
}

double& RHS2D::operator[](std::size_t i)
{
	return v[i];
}


WeightedJacobi2D::WeightedJacobi2D(const PossionMatrix2D& A, const Vector& _f, double _omega) : n(int(sqrt(_f.size())) + 1),
T(SymmetricTridiagonal2D(n, 1 - omega / 4.0 / n / n * A.GetDiag(), -omega / 4.0 / n / n * A.GetAlpha())),
c(omega / A.GetDiag() * _f), omega(_omega) {}

void WeightedJacobi2D::IterateOnce(Vector& v) const
{
	v = T * v + c;
}
